import argparse
import subprocess
import sys

def generate_flamegraph(stack_file, graph_file):
    """Generate a flame graph using the given stack file."""
    flamegraph_cmd = [
        './FlameGraph/flamegraph.pl', 
        '--title="off-cpu flame graph"', 
        stack_file
    ]
    try:
        with open(graph_file, 'w') as out_file:
            subprocess.run(flamegraph_cmd, stdout=out_file, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error generating flame graph: {e}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='Generate off-CPU flame graphs.')
    parser.add_argument('duration', type=int, help='Duration of the trace in seconds.')
    parser.add_argument('-p', '--pid', type=int, help='Trace this PID only.')
    parser.add_argument('-t', '--tid', type=int, help='Trace this TID only.')
    parser.add_argument('-u', '--user-threads-only', action='store_true', help='User threads only.')
    parser.add_argument('-k', '--kernel-threads-only', action='store_true', help='Kernel threads only.')
    parser.add_argument('-U', '--user-stacks-only', action='store_true', help='Show user space stacks only.')
    parser.add_argument('-K', '--kernel-stacks-only', action='store_true', help='Show kernel space stacks only.')
    parser.add_argument('-d', '--delimited', action='store_true', help='Insert delimiter between kernel/user stacks.')
    parser.add_argument('-f', '--folded', action='store_true', help='Output folded format.')
    parser.add_argument('-s', '--offset', action='store_true', help='Show address offsets.')
    parser.add_argument('--stack-storage-size', type=int, help='Number of unique stack traces to store.')
    parser.add_argument('-m', '--min-block-time', type=int, help='Minimum block time in microseconds.')
    parser.add_argument('-M', '--max-block-time', type=int, help='Maximum block time in microseconds.')
    parser.add_argument('--state', type=int, help='Filter on thread state bitmask.')
    parser.add_argument('-o', '--output', type=str, required=True, help='Output file for stack traces.')
    parser.add_argument('-g', '--graph', type=str, required=True, help='Output file for flame graph.')

    args = parser.parse_args()

    # Build the offcputime command
    offcpu_cmd = [
        'python', 'offcputime.py',
    ]
    
    if args.pid:
        offcpu_cmd += ['-p', str(args.pid)]
    if args.tid:
        offcpu_cmd += ['-t', str(args.tid)]
    if args.user_threads_only:
        offcpu_cmd.append('-u')
    if args.kernel_threads_only:
        offcpu_cmd.append('-k')
    if args.user_stacks_only:
        offcpu_cmd.append('-U')
    if args.kernel_stacks_only:
        offcpu_cmd.append('-K')
    if args.delimited:
        offcpu_cmd.append('-d')
    if args.folded:
        offcpu_cmd.append('-f')
    if args.offset:
        offcpu_cmd.append('-s')
    if args.stack_storage_size:
        offcpu_cmd += ['--stack-storage-size', str(args.stack_storage_size)]
    if args.min_block_time:
        offcpu_cmd += ['-m', str(args.min_block_time)]
    if args.max_block_time:
        offcpu_cmd += ['-M', str(args.max_block_time)]
    if args.state:
        offcpu_cmd += ['--state', str(args.state)]
    
    offcpu_cmd += [str(args.duration)]

    # Print the offcputime command for debugging
    print(f"Running command: {' '.join(offcpu_cmd)}")
    
    # Run offcputime and redirect output to specified file
    with open(args.output, 'w') as output_file:
        try:
            subprocess.run(offcpu_cmd, stdout=output_file, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error running offcputime: {e}")
            sys.exit(1)
    
    # Generate flame graph from the stack traces
    generate_flamegraph(args.output, args.graph)

if __name__ == '__main__':
    main()
